ARG IMAGE_NAME
FROM ${IMAGE_NAME}:{{ cuda.version.release_label }}-runtime-{{ cuda.os.distro }}{{ cuda.os.version }}{{ cuda.image_tag_suffix if "image_tag_suffix" in cuda }} as base

{% if "x86_64" in cuda %}
FROM base as base-amd64

ENV NV_CUDA_LIB_VERSION {{ cuda.version.release_label + "-1" }}
ENV NV_NVPROF_VERSION {{ cuda.x86_64.components.nvprof.version }}
ENV NV_CUDA_CUDART_DEV_VERSION {{ cuda.x86_64.components.cudart_dev.version }}
ENV NV_NVML_DEV_VERSION {{ cuda.x86_64.components.nvml_dev.version }}
ENV NV_LIBCUBLAS_DEV_VERSION {{ cuda.x86_64.components.libcublas_dev.version }}
{% if "libnpp_dev" in cuda.x86_64.components %}
    {% set has_libnpp_dev_package = true %}
ENV NV_LIBNPP_DEV_VERSION {{ cuda.x86_64.components.libnpp_dev.version }}
ENV NV_LIBNPP_DEV_PACKAGE libnpp-devel-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBNPP_DEV_VERSION}
{% endif %}
{% if "libnccl2" in cuda.x86_64.components and cuda.x86_64.components.libnccl2 and "libnccl2_dev" in cuda.x86_64.components and cuda.x86_64.components.libnccl2_dev %}
    {% set has_libnccl_dev_package = true %}
ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-devel
ENV NV_LIBNCCL_DEV_PACKAGE_VERSION {{ cuda.x86_64.components.libnccl2_dev.version }}
ENV NCCL_VERSION {{ cuda.x86_64.components.libnccl2_dev.version[:-2] }}
ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}-${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}
{% endif %}

{% endif %}
{% if "arm64" in cuda %}
FROM base as base-arm64

ENV NV_CUDA_LIB_VERSION {{ cuda.version.release_label + "-1" }}
ENV NV_NVPROF_VERSION {{ cuda.arm64.components.nvprof.version }}
ENV NV_CUDA_CUDART_DEV_VERSION {{ cuda.arm64.components.cudart_dev.version }}
ENV NV_NVML_DEV_VERSION {{ cuda.arm64.components.nvml_dev.version }}
ENV NV_LIBCUBLAS_DEV_VERSION {{ cuda.arm64.components.libcublas_dev.version }}
{% if "libnpp_dev" in cuda.arm64.components %}
    {% set has_libnpp_dev_package = true %}
ENV NV_LIBNPP_DEV_VERSION {{ cuda.arm64.components.libnpp_dev.version }}
ENV NV_LIBNPP_DEV_PACKAGE libnpp-devel-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBNPP_DEV_VERSION}
{% endif %}
{% if "libnccl2" in cuda.arm64.components and cuda.arm64.components.libnccl2 and "libnccl2_dev" in cuda.arm64.components and cuda.arm64.components.libnccl2_dev %}
    {% set has_libnccl_dev_package = true %}
ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-devel
ENV NV_LIBNCCL_DEV_PACKAGE_VERSION {{ cuda.arm64.components.libnccl2_dev.version }}
ENV NCCL_VERSION {{ cuda.arm64.components.libnccl2_dev.version[:-2] }}
ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}-${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}
{% endif %}
{% endif %}

{% if "ppc64le" in cuda %}
FROM base as base-ppc64le

ENV NV_CUDA_LIB_VERSION {{ cuda.version.release_label + "-1" }}
ENV NV_NVPROF_VERSION {{ cuda.ppc64le.components.nvprof.version }}
ENV NV_CUDA_CUDART_DEV_VERSION {{ cuda.ppc64le.components.cudart_dev.version }}
ENV NV_NVML_DEV_VERSION {{ cuda.ppc64le.components.nvml_dev.version }}
ENV NV_LIBCUBLAS_DEV_VERSION {{ cuda.ppc64le.components.libcublas_dev.version }}
{% if "libnpp_dev" in cuda.x86_64.components %}
    {% set has_libnpp_dev_package = true %}
ENV NV_LIBNPP_DEV_VERSION {{ cuda.ppc64le.components.libnpp_dev.version }}
ENV NV_LIBNPP_DEV_PACKAGE libnpp-devel-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBNPP_DEV_VERSION}
{% endif %}
{% if "libnccl2" in cuda.ppc64le.components and cuda.ppc64le.components.libnccl2 and "libnccl2_dev" in cuda.ppc64le.components and cuda.ppc64le.components.libnccl2_dev %}
    {% set has_libnccl_dev_package = true %}
ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-devel
ENV NV_LIBNCCL_DEV_PACKAGE_VERSION {{ cuda.ppc64le.components.libnccl2_dev.version }}
ENV NCCL_VERSION {{ cuda.ppc64le.components.libnccl2_dev.version[:-2] }}
ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}-${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda{{ cuda.version.major }}.{{ cuda.version.minor }}
{% endif %}
{% endif %}

FROM base-${TARGETARCH}

ARG TARGETARCH

LABEL maintainer "NVIDIA CORPORATION <sw-cuda-installer@nvidia.com>"

RUN yum install -y \
    make \
    cuda-command-line-tools-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_CUDA_LIB_VERSION} \
    cuda-libraries-devel-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_CUDA_LIB_VERSION} \
    cuda-minimal-build-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_CUDA_LIB_VERSION} \
    cuda-cudart-devel-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_CUDA_CUDART_DEV_VERSION} \
    cuda-nvprof-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_NVPROF_VERSION} \
    cuda-nvml-devel-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_NVML_DEV_VERSION} \
    libcublas-devel-{{ cuda.version.major }}-{{ cuda.version.minor }}-${NV_LIBCUBLAS_DEV_VERSION} \
    {% if has_libnpp_dev_package %}
    ${NV_LIBNPP_DEV_PACKAGE} \
    {% endif %}
    {% if has_libnccl_dev_package %}
    ${NV_LIBNCCL_DEV_PACKAGE} \
    {% endif %}
    && yum clean all \
    && rm -rf /var/cache/yum/*

ENV LIBRARY_PATH /usr/local/cuda/lib64/stubs
